import numpy as np
import matplotlib.pyplot as plt
import os
import sys
from refract import *

freqMin = 1.0
freqMax = 105.0
freqInt = 0.1

freqs = np.arange(freqMin * 1e9, freqMax * 1e9, freqInt * 1e9)

tName = sys.argv[1]

wset1 = sys.argv[2]
wterms1 = wset1.split(",")
wset2 = sys.argv[3]
wterms2 = wset2.split(",")
wset3 = sys.argv[4]
wterms3 = wset3.split(",")

weather1 = [ float(wterms1[0]) + 273.15, float(wterms1[1]) * 100.0, float(wterms1[2]) / 100.0 ]
weather2 = [ float(wterms2[0]) + 273.15, float(wterms2[1]) * 100.0, float(wterms2[2]) / 100.0 ]
weather3 = [ float(wterms3[0]) + 273.15, float(wterms3[1]) * 100.0, float(wterms3[2]) / 100.0 ]

r1 = calcOpacity(freqs, math.radians(90.0), weather1[0], weather1[1], weather1[2])
r2 = calcOpacity(freqs, math.radians(90.0), weather2[0], weather2[1], weather2[2])
r3 = calcOpacity(freqs, math.radians(90.0), weather3[0], weather3[1], weather3[2])

label1 = "T = %.1f K, P = %.1f hPa, H = %.2f" % (weather1[0], weather1[1] / 100.0, weather1[2])
label2 = "T = %.1f K, P = %.1f hPa, H = %.2f" % (weather2[0], weather2[1] / 100.0, weather2[2])
label3 = "T = %.1f K, P = %.1f hPa, H = %.2f" % (weather3[0], weather3[1] / 100.0, weather3[2])

fig = plt.figure(figsize=(11.692, 8.267))
plt.subplot(3, 1, 1)
plt.plot(freqs / 1e9, r1['tau'], "green", label=label1)
plt.plot(freqs / 1e9, r2['tau'], "blue", label=label2)
plt.plot(freqs / 1e9, r3['tau'], "red", label=label3)
plt.title(tName + " opacity")
plt.axis([1.0, 105.0, 0, 30])
plt.legend(fontsize=10)
plt.ylabel("[nepers]")

plt.subplot(3, 1, 2)
plt.plot(freqs / 1e9, r1['Tb'], "green", label=label1)
plt.plot(freqs / 1e9, r2['Tb'], "blue", label=label2)
plt.plot(freqs / 1e9, r3['Tb'], "red", label=label3)
plt.axis([1.0, 105.0, 0, 350])
plt.title(tName + " atmospheric temperature")
plt.ylabel("[K]")

plt.subplot(3, 1, 3)
plt.plot(freqs / 1e9, r1['fac'], "green", label=label1)
plt.plot(freqs / 1e9, r2['fac'], "blue", label=label2)
plt.plot(freqs / 1e9, r3['fac'], "red", label=label3)
plt.axis([1.0, 105.0, 0.0, 1.0])
plt.title(tName + " transmissivity")
plt.xlabel("Frequency [GHz]")
plt.ylabel("[fractional]")

plt.tight_layout()
fname = "atmosphereplots_%s.pdf" % tName
plt.savefig(fname)
